package cn.iocoder.yudao.module.security.service.asset;

import cn.hutool.core.date.DatePattern;
import cn.hutool.core.date.LocalDateTimeUtil;
import cn.iocoder.yudao.framework.common.enums.CommonStatusEnum;
import cn.iocoder.yudao.framework.common.util.date.LocalDateTimeUtils;
import cn.iocoder.yudao.framework.mybatis.core.query.LambdaQueryWrapperX;
import cn.iocoder.yudao.module.security.dal.mysql.asset.SecurityAssetMapper;
import cn.iocoder.yudao.module.security.util.IpUtils;
import cn.iocoder.yudao.module.security.dal.dataobject.asset.SecurityAssetDO;
import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.SortOrder;
import co.elastic.clients.elasticsearch._types.aggregations.AggregationBuilders;
import co.elastic.clients.elasticsearch._types.aggregations.StringTermsBucket;
import co.elastic.clients.elasticsearch._types.aggregations.TermsAggregation;
import co.elastic.clients.elasticsearch.core.SearchRequest;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import co.elastic.clients.util.NamedValue;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import javax.annotation.Resource;
import java.io.IOException;
import java.time.Duration;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;

@Slf4j
@Service
public class SecurityAssetAggService {

    @Resource
    private ElasticsearchClient elasticsearchClient;

    @Resource
    private SecurityAssetMapper assetMapper;

    public int aggAsset() {
        try {
            Long unConfirmCount = assetMapper.selectCount(
                    new LambdaQueryWrapperX<SecurityAssetDO>()
                            .eq(SecurityAssetDO::getDeleted, false)
                            .eq(SecurityAssetDO::getStatus, CommonStatusEnum.DISABLE.getStatus())
            );
            if (unConfirmCount > 100000) {
                log.error("[aggAsset][资产数量超过100000，不进行资产聚合]");
                return 0;
            }
            List<String> list = getSourceIpAgg(1);
            int total = 0;
            //list 1000个一组, 批量处理
            for (int i = 0; i < list.size(); i += 1000) {
                List<String> subList = list.subList(i, Math.min(i + 1000, list.size()));
                List<SecurityAssetDO> securityAssetDOList = new ArrayList<>();

                List<SecurityAssetDO> existList = assetMapper.selectList(
                        new LambdaQueryWrapperX<SecurityAssetDO>()
                                .in(SecurityAssetDO::getIpAddress, subList)
                                .eq(SecurityAssetDO::getDeleted, false)
                );
                Set<String> ipExistSet = existList.stream().map(SecurityAssetDO::getIpAddress).collect(Collectors.toSet());

                for (String ip : subList) {
                    if (ipExistSet.contains(ip)) {
                        continue;
                    }
                    SecurityAssetDO securityAssetDO = new SecurityAssetDO();
                    securityAssetDO.setIpAddress(ip);
                    securityAssetDO.setIpLong(IpUtils.ipv4ToLong(ip));
                    securityAssetDO.setStatus(CommonStatusEnum.DISABLE.getStatus());
                    securityAssetDO.setSourceType("1");
                    securityAssetDOList.add(securityAssetDO);
                }
                total = total + securityAssetDOList.size();
                assetMapper.insertBatch(securityAssetDOList);
            }
            return total;
        } catch (Exception e) {
            log.error("[aggAsset][执行异常]", e);
        }
        return 0;
    }

    public List<String> getSourceIpAgg(int day) {

        LocalDateTime endTime = LocalDateTimeUtil.now();
        LocalDateTime startTime = endTime;
        if (day > 1) {
            startTime = LocalDateTimeUtils.minusTime(Duration.ofDays(day - 1));
        }
        startTime = LocalDateTimeUtil.beginOfDay(startTime);
        String startTimeStr = LocalDateTimeUtil.format(startTime, DatePattern.UTC_SIMPLE_PATTERN) + "+08:00";
        String endTimeStr = LocalDateTimeUtil.format(endTime, DatePattern.UTC_SIMPLE_PATTERN) + "+08:00";

        // 4. 构建聚合分析
        TermsAggregation termsAggregation = AggregationBuilders.terms()
                .field("source.ip")
                .size(100000)
                .order(new NamedValue<SortOrder>("_count", SortOrder.Desc))
                .build();
        SearchRequest request = SearchRequest.of(s -> s
                .index(".ds-logs-pfelk-snort-*")
                //.index("elastalert_log")
                .query(q -> q
                        .bool(b -> b
                                        .filter(f -> f.range(r -> r.date(a -> a.field("@timestamp")
                                                .gte(startTimeStr)
                                                .lte(endTimeStr))))
                                //.filter(f -> f.term(t -> t.field("elastalert.alert").value(0)))
                        )
                )
                .aggregations("source_ip_agg", termsAggregation._toAggregation())
                .size(0)
        );
        List<String> list2 = new ArrayList<>();
        // 执行并处理结果
        try {
            SearchResponse<Object> response = elasticsearchClient.search(request, Object.class);

            // 7. 解析聚合结果
            List<StringTermsBucket> buckets = response.aggregations()
                    .get("source_ip_agg")
                    .sterms()
                    .buckets()
                    .array();
            for (StringTermsBucket bucket : buckets) {
                String key = bucket.key().stringValue();
                if (IpUtils.isIpv4Address(key) && IpUtils.isInternalIP(key)) {
                    list2.add(key);
                }
            }
        } catch (IOException e) {
            log.error("elasticsearch request err,request: {}", request, e);
        }

        return list2;
    }

}
